#pragma once

#include <vector>
#include <algorithm> 
#include <functional> 
#include <windows.h>
#include "disasm.h"

#include "Common.h"

#define IS_FUNC_END( x ) ( x == 0xc2 || x == 0xc3 )

typedef struct _BLOCK_FUNC
{
	ULONG uBlockAddr;
	ULONG uBlockLen;

}BLOCK_FUNC, *PBLOCK_FUNC;


typedef std::vector <ULONG> Addr_List;

typedef std::pair<ULONG, ULONG> SECTION_LIMIT;

class CRetFinder
{
public:
	CRetFinder(void);
	CRetFinder(ULONG uBaseAddr, ULONG uParamSize, ULONG uFuncLen);
	virtual ~CRetFinder(void);
public:
	VOID SetBaseAddr( ULONG uBaseAddr );
	VOID SetParamSize( ULONG uParamSize );
	BOOL GetRetOffset( PULONG pOffset );
	VOID SetFunctionSize( ULONG uFuncSize );
	DWORD GetSpcifiedCodeOffset( PBYTE pCode, int uSize );
	VOID InitDisasm();
	BOOL FindBlock();
	VOID PrintBlock();
    VOID PrintBlock(std::string& OutBuffer);
private:
	BOOL IsLoopAddrInBlock( ULONG uLoopAddr );
	BOOL GetUnLoopList( Addr_List& AddList );
	BOOL FindBlock( ULONG uDiasmAddr, BOOL bIsRollUp, Addr_List& AddList );
	BOOL GetRetOffset( ULONG uDasmAddr, PULONG pOffset );
	VOID AddBlock2List( PBLOCK_FUNC pBlock );
	VOID AddLoopAddr2List( Addr_List& AddList, ULONG uAddr );
	VOID AddNewBlock( ULONG StartAddr, ULONG EndAddr );
    BOOL IsJmpNext( t_disasm* da, ULONG uLen );
    BOOL IsInLimitSection(ULONG Address);
private:
	ULONG m_BaseAddr;
	ULONG m_ParamSize;
	Addr_List m_LoopAddressList;
	std::vector <PBLOCK_FUNC> m_BlockList;
    SECTION_LIMIT m_SectionLimit;
	BOOL m_Ideal; 
	BOOL m_LowerCase; 
	BOOL m_PutDefseg;
	ULONG m_FuncLen;
    
public:
    ULONG GetBlockCount(void);
    VOID GetBlockBuffer(LPVOID Buffer, ULONG Length);
    VOID SetSectionLimit( ULONG SectionStart, ULONG SectionLength );
    
};
